/*
 *  Copyright 2016 The WebRTC Project Authors. All rights reserved.
 *
 *  Use of this source code is governed by a BSD-style license
 *  that can be found in the LICENSE file in the root of the source
 *  tree. An additional intellectual property rights grant can be found
 *  in the file PATENTS.  All contributing project authors may
 *  be found in the AUTHORS file in the root of the source tree.
 */

package org.webrtc;

import android.os.Handler;
import android.os.HandlerThread;

import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.charset.Charset;
import java.util.concurrent.CountDownLatch;

/**
 * Can be used to save the video frames to file.
 */
public class VideoFileRenderer implements VideoSink {
    private static final String TAG = "VideoFileRenderer";

    private final HandlerThread renderThread;
    private final Handler renderThreadHandler;
    private final HandlerThread fileThread;
    private final Handler fileThreadHandler;
    private final FileOutputStream videoOutFile;
    private final String outputFileName;
    private final int outputFileWidth;
    private final int outputFileHeight;
    private final int outputFrameSize;
    private final ByteBuffer outputFrameBuffer;
    private EglBase eglBase;
    private YuvConverter yuvConverter;
    private int frameCount;

    public VideoFileRenderer(String outputFile, int outputFileWidth, int outputFileHeight,
                             final EglBase.Context sharedContext) throws IOException {
        if ((outputFileWidth % 2) == 1 || (outputFileHeight % 2) == 1) {
            throw new IllegalArgumentException("Does not support uneven width or height");
        }

        this.outputFileName = outputFile;
        this.outputFileWidth = outputFileWidth;
        this.outputFileHeight = outputFileHeight;

        outputFrameSize = outputFileWidth * outputFileHeight * 3 / 2;
        outputFrameBuffer = ByteBuffer.allocateDirect(outputFrameSize);

        videoOutFile = new FileOutputStream(outputFile);
        videoOutFile.write(
                ("YUV4MPEG2 C420 W" + outputFileWidth + " H" + outputFileHeight + " Ip F30:1 A1:1\n")
                        .getBytes(Charset.forName("US-ASCII")));

        renderThread = new HandlerThread(TAG + "RenderThread");
        renderThread.start();
        renderThreadHandler = new Handler(renderThread.getLooper());

        fileThread = new HandlerThread(TAG + "FileThread");
        fileThread.start();
        fileThreadHandler = new Handler(fileThread.getLooper());

        ThreadUtils.invokeAtFrontUninterruptibly(renderThreadHandler, new Runnable() {
            @Override
            public void run() {
                eglBase = EglBase.create(sharedContext, EglBase.CONFIG_PIXEL_BUFFER);
                eglBase.createDummyPbufferSurface();
                eglBase.makeCurrent();
                yuvConverter = new YuvConverter();
            }
        });
    }

    @Override
    public void onFrame(VideoFrame frame) {
        frame.retain();
        renderThreadHandler.post(() -> renderFrameOnRenderThread(frame));
    }

    private void renderFrameOnRenderThread(VideoFrame frame) {
        final VideoFrame.Buffer buffer = frame.getBuffer();

        // If the frame is rotated, it will be applied after cropAndScale. Therefore, if the frame is
        // rotated by 90 degrees, swap width and height.
        final int targetWidth = frame.getRotation() % 180 == 0 ? outputFileWidth : outputFileHeight;
        final int targetHeight = frame.getRotation() % 180 == 0 ? outputFileHeight : outputFileWidth;

        final float frameAspectRatio = (float) buffer.getWidth() / (float) buffer.getHeight();
        final float fileAspectRatio = (float) targetWidth / (float) targetHeight;

        // Calculate cropping to equalize the aspect ratio.
        int cropWidth = buffer.getWidth();
        int cropHeight = buffer.getHeight();
        if (fileAspectRatio > frameAspectRatio) {
            cropHeight = (int) (cropHeight * (frameAspectRatio / fileAspectRatio));
        } else {
            cropWidth = (int) (cropWidth * (fileAspectRatio / frameAspectRatio));
        }

        final int cropX = (buffer.getWidth() - cropWidth) / 2;
        final int cropY = (buffer.getHeight() - cropHeight) / 2;

        final VideoFrame.Buffer scaledBuffer =
                buffer.cropAndScale(cropX, cropY, cropWidth, cropHeight, targetWidth, targetHeight);
        frame.release();

        final VideoFrame.I420Buffer i420 = scaledBuffer.toI420();
        scaledBuffer.release();

        fileThreadHandler.post(() -> {
            YuvHelper.I420Rotate(i420.getDataY(), i420.getStrideY(), i420.getDataU(), i420.getStrideU(),
                    i420.getDataV(), i420.getStrideV(), outputFrameBuffer, i420.getWidth(), i420.getHeight(),
                    frame.getRotation());
            i420.release();

            try {
                videoOutFile.write("FRAME\n".getBytes(Charset.forName("US-ASCII")));
                videoOutFile.write(
                        outputFrameBuffer.array(), outputFrameBuffer.arrayOffset(), outputFrameSize);
            } catch (IOException e) {
                throw new RuntimeException("Error writing video to disk", e);
            }
            frameCount++;
        });
    }

    /**
     * Release all resources. All already posted frames will be rendered first.
     */
    public void release() {
        final CountDownLatch cleanupBarrier = new CountDownLatch(1);
        renderThreadHandler.post(() -> {
            yuvConverter.release();
            eglBase.release();
            renderThread.quit();
            cleanupBarrier.countDown();
        });
        ThreadUtils.awaitUninterruptibly(cleanupBarrier);
        fileThreadHandler.post(() -> {
            try {
                videoOutFile.close();
                Logging.d(TAG,
                        "Video written to disk as " + outputFileName + ". The number of frames is " + frameCount
                                + " and the dimensions of the frames are " + outputFileWidth + "x"
                                + outputFileHeight + ".");
            } catch (IOException e) {
                throw new RuntimeException("Error closing output file", e);
            }
            fileThread.quit();
        });
        try {
            fileThread.join();
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            Logging.e(TAG, "Interrupted while waiting for the write to disk to complete.", e);
        }
    }
}
